import os
import SimpleITK as sitk
import numpy as np
import cv2
import argparse
from tqdm import tqdm
from multiprocessing import Pool, cpu_count

def process_single_volume(args):
    filename, input_dir, output_dir = args
    image_path = os.path.join(input_dir, "image", filename)
    label_path = os.path.join(input_dir, "label", filename)
    
    # 检查label文件是否存在
    if not os.path.exists(label_path):
        print(f"Warning: Label file not found for {filename}, skipping...")
        return
    
    try:
        # 读取图像和标注数据
        image_vol = sitk.ReadImage(image_path)
        label_vol = sitk.ReadImage(label_path)
        image_arr:np.ndarray = sitk.GetArrayFromImage(image_vol)
        label_arr:np.ndarray = sitk.GetArrayFromImage(label_vol)
        
        # 找到包含标注的切面
        valid_slices = np.any(label_arr != 0, axis=(1,2))
        valid_image_arr = image_arr[valid_slices]
        valid_label_arr = label_arr[valid_slices]
        
        if valid_slices.any() is False:
            print(f"Warning: No labeled slice found in {filename}, skipping...")
            return
        
        # 获取包含标注的切面索引
        for i, (img, ann) in enumerate(zip(valid_image_arr, valid_label_arr)):
            output_name = f"{os.path.splitext(filename)[0]}_{i}.tiff"
            img_output_dir = os.path.join(output_dir, "image")
            lbl_output_dir = os.path.join(output_dir, "label")
            os.makedirs(img_output_dir, exist_ok=True)
            os.makedirs(lbl_output_dir, exist_ok=True)
            cv2.imwrite(os.path.join(img_output_dir, output_name),
                        img.astype(np.int16))
            cv2.imwrite(os.path.join(lbl_output_dir, output_name),
                        ann.astype(np.uint8))
        
    except Exception as e:
        print(f"Error processing {filename}: {str(e)}")
        return

def convert_volume_to_2d_samples(input_dir, output_dir):
    # 确保输入目录存在
    if not os.path.exists(input_dir) or not os.path.exists(os.path.join(input_dir, "image")):
        raise ValueError(f"Input directory {input_dir} or its 'image' subdirectory does not exist")
    
    # 获取所有.mha文件
    image_files = [f for f in os.listdir(os.path.join(input_dir, "image")) if f.endswith(".mha")]
    if not image_files:
        print("Warning: No .mha files found in input directory")
        return
    
    # 准备进程池
    num_processes = max(1, cpu_count() - 1)
    pool = Pool(processes=num_processes)
    
    # 构建参数列表
    args_list = [(f, input_dir, output_dir) for f in image_files]
    
    # 使用进程池处理并显示进度条
    list(tqdm(pool.imap(process_single_volume, args_list), 
             total=len(image_files), 
             desc="Processing volumes"))
    
    pool.close()
    pool.join()

def parse_args():
    parser = argparse.ArgumentParser(description="Convert 3D volumes to 2D samples")
    parser.add_argument("input_dir", help="Path to input folder with 'image' and 'label' subfolders")
    parser.add_argument("output_dir", help="Path where 2D slices will be stored")
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    convert_volume_to_2d_samples(args.input_dir, args.output_dir)